package geppetto.phraseHMM;



import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentEvaluator;
import geppetto.cat.alignments.AlignmentStats;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.alignments.AlignmentEvaluator.Evaluation;
import geppetto.cat.common.MyArrays;
import geppetto.cat.common.StaticTools;
import geppetto.cat.constrains.ConstrainedProjectionStats;
import geppetto.cat.constrains.SentenceConstrainedProjectionStats;
import geppetto.cat.constrains.SymmetryConstrains;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractDoubleDistortionTable;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.AbstractSparseTranslationTable;
import geppetto.cat.models.stats.EStepStats;
import geppetto.cat.models.stats.SentenceEStepStats;

import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.util.ArrayList;

public class SymmetricHMM extends AbstractModel{

	public RegularHMM forward;
	public RegularHMM backward;
	public double optimizationPrecision;
	public double slack;
	public int _maxStepSize; 
	public int _maxNumberIterations; 
	
	public SymmetricHMM(BilingualCorpus corpus,BilingualCorpus revCorpus, 
			double smoothing ,
			double epsilon, double slack, int maxStepSize, int maxNumberProjectionIterations) {
		this(corpus,revCorpus, null,null, smoothing,epsilon,slack,maxStepSize,maxNumberProjectionIterations);
	}

	public SymmetricHMM(BilingualCorpus corpus, BilingualCorpus revCorpus, AbstractSparseTranslationTable tt,
			AbstractSparseTranslationTable revtt, 
			double smoothing ,
			double epsilon, double slack, int maxStepSize, int maxNumberProjectionIterations) {
		this(corpus, revCorpus, tt, revtt, null,null,smoothing, epsilon,slack,maxStepSize,maxNumberProjectionIterations);
	}

	public SymmetricHMM(BilingualCorpus corpus, BilingualCorpus revCorpus, AbstractSparseTranslationTable tt,
			AbstractSparseTranslationTable revtt,
			AbstractDoubleDistortionTable distortion,AbstractDoubleDistortionTable revDistortion,
			double smoothing ,
			double epsilon, double slack, int maxStepSize, int maxNumberProjectionIterations) {
		forward = new RegularHMM(corpus,tt,distortion,smoothing);
		backward = new RegularHMM(revCorpus,revtt,revDistortion,smoothing);
		_corpus = corpus;
		this.optimizationPrecision = epsilon;
		this.slack = slack;
		_maxStepSize = maxStepSize;
		_maxNumberIterations = maxNumberProjectionIterations;
	}

	public void printStamp(PrintStream file){
		super.printStamp(file);
		file.println("Epsilon: " + optimizationPrecision);
		file.println("Slack: " + slack);
		file.println("Max Step Size: " + _maxStepSize);
		file.println("Max Number of Projection : " + _maxNumberIterations);
	}
	
	//Load and Save File
//	 /////////////////SAVE AND LOAD MODELS

	
	SymmetricHMM(RegularHMM forward, RegularHMM backward) {
		this.forward = forward;
		this.backward = backward;
		this._corpus = forward._corpus;
	}

	public SymmetricHMM getBackwardModel() {
		SymmetricHMM res = new SymmetricHMM(backward, forward);
		res._trained = _trained;
		res._maxNumberIterations = _maxNumberIterations;
		res._maxStepSize = _maxStepSize;
		res.slack = slack;
		res.optimizationPrecision = optimizationPrecision;
		return res;
	}
	
	public void saveModel(String directory) {
		StaticTools.createDir(directory + "/forward/");
		StaticTools.createDir(directory + "/backward/");
		forward.saveModel(directory + "/forward/");
		backward.saveModel(directory + "/backward/");
		try {
			PrintStream file = new PrintStream(new FileOutputStream(directory+ "stamp"));
			printStamp(file);
		} catch (FileNotFoundException e) {
			System.out.println(getName()+": Could not save model");
			System.exit(-1);
		}
	}

	
	public static SymmetricHMM loadModel(BilingualCorpus forwardCorpus,
			BilingualCorpus backwardCorpus, String directory,
			double epsilon, double slack, int maxSteps, int maxIterations) {
		SymmetricHMM model = new SymmetricHMM();
		try {
			System.out.println("Reading Forward: " + directory + "forward");
			model.forward = RegularHMM.loadModel(forwardCorpus, directory + "forward");
			model._corpus = forwardCorpus;
			System.out.println("Reading Forward: " + directory + "backward");
			model.backward = RegularHMM.loadModel(backwardCorpus, directory+ "backward");
		} catch (NumberFormatException e) {
			System.out
					.println(model.forward.getName() + "Could not read model number conversion exception"
							+ e.toString());
			System.exit(-1);

		} 
		model.optimizationPrecision=epsilon;
		model.slack=slack;
		model._maxStepSize = maxSteps;
		model._maxNumberIterations = maxIterations;
		return model;
	}
	
	
	SymmetricHMM(){
			
	}
	public String getName() {
		return "Symmetric Constrain HMM";
	}

	
	public SentenceConstrainedProjectionStats projectPosteriors(){
		SymmetryConstrains constraints = new SymmetryConstrains(forward, backward, optimizationPrecision, slack,_maxStepSize,_maxNumberIterations);
//		SentenceConstrainedProjectionStats stats = constraints.steepestAscentProjection();
		SentenceConstrainedProjectionStats stats = constraints.conjugateGradientProjection2();
	//	System.out.println(stats.prettyPrint());
		return stats;
	}
	
	public EStepStats createModelStats(){
		ConstrainedProjectionStats pstats = new ConstrainedProjectionStats();
		EStepStats stats = new EStepStats();
		stats.pstats = pstats;
		return stats;
	}
	
	public SentenceEStepStats sentenceEStep(){
		SentenceEStepStats stats = new SentenceEStepStats();
		forward.makeCaches();
		forward.makePosteriors(forward._sSize, forward._fSize);
		backward.makeCaches();
		backward.makePosteriors(backward._sSize, backward._fSize);
		stats.loglikelihood = (forward.getRealLikelihood()+backward.getRealLikelihood())/2;		
		if(stats.loglikelihood < 1.0e-200){
			System.out.println("Skipping a sentence due to likelihood ");
			return stats;
		}
		SentenceConstrainedProjectionStats pstats = projectPosteriors();
		stats.pstats = pstats;
		return stats;
	}
	
	
	
	public EStepStats eStep() {
		
		forward._currentTrainingIteration = _currentTrainingIteration;
		backward._currentTrainingIteration = _currentTrainingIteration;
		forward.initializeCounts();
		backward.initializeCounts();
		EStepStats stats = createModelStats();
		for (int i = 0; i < forward._nSentences; i++) {
			forward.initSentence(i, BilingualCorpus.TRAIN_CORPUS);
			backward.initSentence(i, BilingualCorpus.TRAIN_CORPUS);
			final int sSize = forward._sourceSentenceIDS.length;
			final int fSize = forward._foreignSentenceIDS.length;
			SentenceEStepStats sstats = sentenceEStep();
			stats.add(sstats,sSize);	
			if(sstats.loglikelihood < 1.0e-200){
				forward._previousSSize = sSize;
				backward._previousSSize = fSize;
				continue;
			}	
			forward.updateFractionalCounts();
			backward.updateFractionalCounts();
			forward._previousSSize = sSize;
			backward._previousSSize = fSize;
		}
		return stats;
	}


	public void finalizeTrain() {
		forward.finalizeTrain();
		backward.finalizeTrain();
	}

	
	public double getNullPhrasePosterior(int sentenceNumber, byte sentenceSource, int[] foreingSentence, int[] sourceSentence, int startForeignIndex, int endForeignIndex) {
		return forward.getNullPhrasePosterior(sentenceNumber, sentenceSource, foreingSentence, sourceSentence, startForeignIndex, endForeignIndex);
	}

	
	public double getPhrasePosterior(int sentenceNumber, byte sentenceSource, int[] foreingSentence, int[] sourceSentence, int startSourceIndex, int endSourceIndex, int startForeignIndex, int endForeignIndex) {
		return forward.getPhrasePosterior(sentenceNumber, sentenceSource, foreingSentence, sourceSentence, startSourceIndex, endSourceIndex, startForeignIndex, endForeignIndex);
	}


	public void initializeTrain() {
		forward.initializeTrain();
		backward.initializeTrain();
		forward._trained = true;
		backward._trained = true;
		_trained=true;
		
	}


	public void mStep() {
		forward.mStep();
		backward.mStep();
		
	}
	
	/**
	 * We are paying the prince twice by calling this to the forward and backward models....
	 */
	public Alignment posteriorDecodingAlignment(int sentenceNumber, byte sentenceSource, float treshhold,boolean projectPosteriors, ConstrainedProjectionStats stats) {
		forward.initDecoding(sentenceNumber, sentenceSource);
		backward.initDecoding(sentenceNumber, sentenceSource);
		if(projectPosteriors){
			SentenceConstrainedProjectionStats sstats = projectPosteriors();
			if(stats != null){
				stats.add(sstats);
			}
		}	
		
	//	forward._statePosteriors.print(forward._fSize, forward._sSize*2);
		Alignment a = new Alignment(sentenceNumber, sentenceSource, forward._sSize,forward._fSize);
		for (int fIndex = 0; fIndex < forward._fSize; fIndex++) {
			for (int sIndex = 0; sIndex < forward._sSize; sIndex++) {
				if (forward._statePosteriors.getProb(fIndex, sIndex) > treshhold) {
					a.add(sIndex, fIndex);
				}
				a.addPosterior(sIndex, fIndex, forward._statePosteriors.getProb(
						fIndex, sIndex));
			}
		}
		return a;
		
		
	}

	public Alignment viterbiAlignment(int sentenceNumber, byte sentenceSource, boolean projectPosteriors, ConstrainedProjectionStats stats) {
		forward.initDecoding(sentenceNumber, sentenceSource);
		backward.initDecoding(sentenceNumber, sentenceSource);
		if(projectPosteriors){
			stats.add(projectPosteriors());
		}
		Trellis viterbi = new Trellis(forward._fSize, forward._sSize * 2, true);
		int[][] states = new int[forward._fSize][forward._sSize * 2];
		for (int i = 1; i < forward._fSize; i++) {
			java.util.Arrays.fill(states[i], -1);
		}
		for (int sourceIndex = 0; sourceIndex < forward._sSize * 2; sourceIndex++) {
			double prob = forward.getInitProb(sourceIndex, forward._sSize)
					* forward.getObservationProbability(0, sourceIndex);
			viterbi.setProb(0, sourceIndex, prob);
			states[0][sourceIndex] = 0;
		}
		for (int pos = 1; pos < forward._fSize; pos++) {
			for (int sourceIndex = 0; sourceIndex < forward._sSize * 2; sourceIndex++) {
				double observation = forward.getObservationProbability(pos, sourceIndex);
				double prob = 0;
				double max = -1;
				int maxState = -2;
				for (int prevSourceIndex = 0; prevSourceIndex < forward._sSize * 2; prevSourceIndex++) {
					double viter = viterbi.getProb(pos - 1, prevSourceIndex);
					double dist = forward.getTransitionProbability(prevSourceIndex,
							sourceIndex, forward._sSize, forward._fSize);
					prob = viter * dist;
					if (prob > max) {
						max = prob;
						maxState = prevSourceIndex;
					}

				}
				viterbi.setProb(pos, sourceIndex, max * observation);
				states[pos][sourceIndex] = maxState;
			}
		}
//		 viterbi.print(_fSize,_sSize);
//		 cat.common.StaticTools.printMatrix(states, _fSize, _sSize, "states", System.out);

		int[] viterbiPath = new int[forward._fSize];
		java.util.Arrays.fill(viterbiPath, -2);

		double max = -1;
		for (int sourceIndex = 0; sourceIndex < forward._sSize * 2; sourceIndex++) {
			double prob = viterbi.getProb(forward._fSize - 1, sourceIndex);
			if (prob > max) {
				max = prob;
				viterbiPath[forward._fSize - 1] = sourceIndex;
			}
		}

		// System.out.println(fSize);
		for (int pos = forward._fSize - 2; pos >= 0; pos--) {
			int currentState = viterbiPath[pos + 1];
			//Debug code
			if(currentState == -2){
				System.out.println("Sentence  " + forward._sentenceNumber + " sentence source" + forward._sentenceSource );
				System.out.println("Fsize " + forward._fSize + " ssize " + forward._sSize );
				System.out.println("Corpus max Fsize " + _corpus.getMaxForeignLen() + " ssize " + _corpus.getMaxSourceLen());
				forward.printForwardBackwardTables();
				forward.printPosteriors();
				System.out.println("states");
				 MyArrays.printIntArray(states, "states", forward._fSize, forward._sSize*2);
				//System.out.println("viterbi");
				//viterbi.print(_fSize, _sSize*2);
				throw new RuntimeException("Viterbi decoding, path as a -2 value");
			}
			viterbiPath[pos] = states[pos + 1][currentState];
		}

		 //MyArrays.printIntArray(viterbiPath,"ViterbiPath");

		Alignment a = new Alignment(sentenceNumber, sentenceSource, forward._sSize,
				forward._fSize);
		for (int i = 0; i < viterbiPath.length; i++) {
			if (viterbiPath[i] < forward._sSize) {
				a.add(viterbiPath[i], i); // Caso em que nao e null
				
			}
		}
		
		//Add posteriors for all positions
		for(int i = 0; i < forward._sSize; i++){
			for(int j = 0; j < forward._fSize; j++){
				a.addPosterior(i, j, forward._statePosteriors.getProb(j,i));
			}
		}
		return a;
	}

//	public AlignmentsSet viterbiAlignments(byte sentenceSource,boolean projectPosteriors) {
//		
//		return forward.viterbiAlignments(sentenceSource);
//	}
	
	public static  Trellis posteriorDifferences(Trellis forward, Trellis backward,int sSize, int fSize){
		Trellis result = new Trellis(fSize,sSize);
		for (int state = 0; state < sSize; state++) {
			for (int pos = 0; pos < fSize; pos++) {
				result.setProb(pos, state, forward.getProb(pos, state)- backward.getProb(state, pos));
			}
		}
		return result;
	}
	
	public static void main(String[] args) throws IOException {
		String corpusDescription = args[0];
		int size = Integer.parseInt(args[1]); 
		int maxSentenceSize = Integer.parseInt(args[2]); 
		int numberIterations = Integer.parseInt(args[3]); 
		double smoothing = Double.parseDouble(args[4]); 
		double slack = Double.parseDouble(args[5]); 
		double epsilon = Double.parseDouble(args[6]); 
		int maxStepSize = Integer.parseInt(args[7]);
		int maxNumberOfProjectionIterations = Integer.parseInt(args[8]);
		String baseOutput = args[9];
		String outputFile = args[10];
		boolean trainWithResults = Boolean.parseBoolean(args[11]);
		int numberIterationsWithResults = Integer.parseInt(args[12]);
		boolean saveModel = Boolean.parseBoolean(args[13]);
		String saveModelDir = args[14];
		boolean useSymmetricM1 = Boolean.parseBoolean(args[15]);
		boolean useBijectiveHMM = Boolean.parseBoolean(args[16]);
		
		// System.out.println("Corpus "+corpusName);
		System.out.println("Size " + size);
		System.out.println("Max Sentence size " + maxSentenceSize);
		System.out.println("Number of iterations " + numberIterations);
		System.out.println("smoothing " + smoothing);
		System.out.println("Slack " + slack);
		System.out.println("Epsilon " + epsilon);
		System.out.println("Max Step Size" + maxStepSize);
		System.out.println("Max Number of Iterations" + maxNumberOfProjectionIterations);
		System.out.println("BaseOutpup " + baseOutput);
		System.out.println("OutputFile " + outputFile);
		System.out.println("Train with results " + trainWithResults);
		System.out.println("Number of Iterations with results "
				+ numberIterationsWithResults);
		System.out.println("Saving model " + saveModel + " to " + saveModelDir);
		System.out.println("Use Symmetric M1 " + useSymmetricM1);
		System.out.println("Use Bijective HMM " + useBijectiveHMM);
		BilingualCorpus forwardCorpus = BilingualCorpus.getCorpusFromFileDescription(
				corpusDescription, size, maxSentenceSize);
		BilingualCorpus backwardCorpus = forwardCorpus.reverse();
		IBMM1 m1;
		IBMM1 m1Back;
		if(useSymmetricM1){
			SymmetricM1 sm1 = new SymmetricM1(forwardCorpus,backwardCorpus,smoothing,epsilon,slack,maxStepSize,maxNumberOfProjectionIterations);
			sm1.train(numberIterations,false,"");
			m1 = sm1.forward;
			m1Back = sm1.backward;
		}else {
			m1 = new IBMM1(forwardCorpus,smoothing);
			m1Back = new IBMM1(backwardCorpus,smoothing);
			m1.train(numberIterations,false,"");
			m1Back.train(numberIterations,false,"");
		}
		String baseDir = saveModelDir + "/symmetric/" + forwardCorpus.getName() + "/" + size + "/";
		if(saveModel){
			StaticTools.createDir(baseDir);
		}
		SymmetricHMM mhmm;
		if(useBijectiveHMM){
			BijectiveHMM bf = new BijectiveHMM(forwardCorpus,m1._tb,0,0.005,0,200,200);
			bf.train(numberIterations,false,"");
			BijectiveHMM bb = new BijectiveHMM(backwardCorpus,m1Back._tb,0,0.005,0,200,200);
			bb.train(numberIterations,false,"");
			 mhmm = new SymmetricHMM(forwardCorpus, backwardCorpus,bf._tb,bb._tb,smoothing,epsilon,slack,maxSentenceSize,maxNumberOfProjectionIterations);
		}else{
			 mhmm = new SymmetricHMM(forwardCorpus, backwardCorpus,m1._tb,m1Back._tb,smoothing,epsilon,slack,maxSentenceSize,maxNumberOfProjectionIterations);
		}
		
		//SymmetricHMM mhmm = new SymmetricHMM(forwardCorpus, backwardCorpus,m1._tb,m1Back._tb,smoothing,epsilon,slack,maxSentenceSize,maxNumberOfProjectionIterations);		
		if (!trainWithResults) {
			mhmm.train(numberIterations,saveModel,baseDir);
		} else {
			ArrayList<Evaluation[]> evalsList = mhmm.trainWithResults(
					numberIterationsWithResults, BilingualCorpus.DEV_CORPUS,saveModel,baseDir);
			System.out.println("Viterbi all");
			Evaluation[] evals = evalsList.get(0);
			for (int i = 0; i < evals.length; i++) {
				System.out.println("Iter " + i + evals[i]);
			}
			System.out.println(" ----- ");
			System.out.println("Viterbi Rare");
			evals = evalsList.get(1);
			for (int i = 0; i < evals.length; i++) {
				System.out.println("Iter " + i + evals[i]);
			}
			System.out.println(" ----- ");
			System.out.println("Viterbi Common");
			evals = evalsList.get(2);
			for (int i = 0; i < evals.length; i++) {
				System.out.println("Iter " + i + evals[i]);
			}
			System.out.println(" ----- ");
			
			System.out.println("Precision all");
			 evals = evalsList.get(3);
			for (int i = 0; i < evals.length; i++) {
				System.out.println("Iter " + i + evals[i]);
			}
			System.out.println(" ----- ");
			System.out.println("Precision Rare");
			evals = evalsList.get(4);
			for (int i = 0; i < evals.length; i++) {
				System.out.println("Iter " + i + evals[i]);
			}
			System.out.println(" ----- ");
			System.out.println("Precision Common");
			evals = evalsList.get(5);
			for (int i = 0; i < evals.length; i++) {
				System.out.println("Iter " + i + evals[i]);
			}
			System.out.println(" ----- ");
		}

		
		System.out.println("Done with training");
		
		
		
		if(saveModel){
			StaticTools.createDir(baseDir);
			mhmm.saveModel(baseDir);
		}
		
		AlignmentsSet sa = mhmm.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
		AlignmentsSet gold = forwardCorpus.getGold();
		System.out.println("Gold size" + gold.size() + " viterbi size "
				+ sa.size());
		Evaluation eval2 = AlignmentEvaluator.evaluate(sa, forwardCorpus.getGold());
		AlignmentStats.printPhrases(sa, true, false);
		System.out.println("Segment viterbi " + eval2);
//		mhmm.outputAlignments(sa, gold, corpus, new PrintStream(baseOutput
//				+ "/latex/" + outputFile + "-viter.tex"));
		
		
		// Viterbi with projection
		AlignmentsSet saP = mhmm.viterbiAlignments(BilingualCorpus.TEST_CORPUS,true);
		Evaluation eval2P = AlignmentEvaluator.evaluate(saP, forwardCorpus.getGold());
		AlignmentStats.printPhrases(saP, true, false);
		System.out.println("Segment viterbi with projection " + eval2P);
//		mhmm.outputAlignments(sa, gold, corpus, new PrintStream(baseOutput
//				+ "/latex/" + outputFile + "-viter.tex"));

		float treshold = mhmm.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,false);
		
		AlignmentsSet sa2 = mhmm.posteriorAlignments(
				BilingualCorpus.TEST_CORPUS, treshold,false,false);
		Evaluation eval22 = AlignmentEvaluator.evaluate(sa2, forwardCorpus.getGold());
		// sa2.outputWithStatistics(corpus, System.out);
		AlignmentStats.printPhrases(sa2, true, false);
		System.out.println("Posterioir decoding " + eval22);
//		mhmm.outputAlignments(sa2, gold, corpus, new PrintStream(baseOutput
//				+ "/latex/" + outputFile + "-post.tex"));
		
		float tresholdP = mhmm.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,true);
		AlignmentsSet sa2P = mhmm.posteriorAlignments(
				BilingualCorpus.TEST_CORPUS, tresholdP,true,false);
		Evaluation eval22P = AlignmentEvaluator.evaluate(sa2P, forwardCorpus.getGold());
		// sa2.outputWithStatistics(corpus, System.out);
		AlignmentStats.printPhrases(sa2P, true, false);
		System.out.println("Posterioir decoding with projection" + eval22P);
//		mhmm.outputAlignments(sa2, gold, corpus, new PrintStream(baseOutput
//				+ "/latex/" + outputFile + "-post.tex"));

		
		System.out.println("Recall curves");
		mhmm.printPosteriorCurves(forwardCorpus, BilingualCorpus.TEST_CORPUS, true, "Forward-P",100);
		

	}

	
	
}
